#!/usr/bin/env python3
"""
增强版无人机导航系统模拟（带 D*-Lite 支持未知 / 动态障碍）
将部分障碍设置为动态，并模拟无人机带有有限传感器逐步发现障碍后触发重规划。
"""

import numpy as np
import matplotlib.pyplot as plt
from matplotlib.animation import FuncAnimation
import math
import time
from queue import PriorityQueue
import heapq
import logging
from logging.handlers import RotatingFileHandler
import csv
import os
from datetime import datetime
from mpl_toolkits.mplot3d.art3d import Poly3DCollection
import copy
import random

plt.rcParams['font.sans-serif'] = ['SimHei']
plt.rcParams['axes.unicode_minus'] = False

# ---------- Logging & DataLogger (和原来保持一致) ----------
def setup_logging():
    if not os.path.exists('logs'):
        os.makedirs('logs')
    log_filename = f"logs/drone_navigation_{datetime.now().strftime('%Y%m%d_%H%M%S')}.log"
    log_format = '%(asctime)s - %(levelname)s - %(message)s'
    logging.basicConfig(
        level=logging.INFO,
        format=log_format,
        handlers=[
            RotatingFileHandler(log_filename, maxBytes=10 * 1024 * 1024, backupCount=5),
            logging.StreamHandler()
        ]
    )
    return log_filename

class DataLogger:
    def __init__(self):
        if not os.path.exists('data'):
            os.makedirs('data')
        self.filename = f"data/drone_data_{datetime.now().strftime('%Y%m%d_%H%M%S')}.csv"
        with open(self.filename, 'w', newline='') as csvfile:
            writer = csv.writer(csvfile)
            writer.writerow([
                'timestamp', 'pos_x', 'pos_y', 'pos_z',
                'vel_x', 'vel_y', 'vel_z',
                'target_x', 'target_y', 'target_z',
                'error_x', 'error_y', 'error_z', 'total_error',
                'control_force_x', 'control_force_y', 'control_force_z'
            ])
    def log_data(self, timestamp, position, velocity, target, control_force):
        error = target - position
        total_error = np.linalg.norm(error)
        with open(self.filename, 'a', newline='') as csvfile:
            writer = csv.writer(csvfile)
            writer.writerow([
                timestamp,
                position[0], position[1], position[2],
                velocity[0], velocity[1], velocity[2],
                target[0], target[1], target[2],
                error[0], error[1], error[2], total_error,
                control_force[0], control_force[1], control_force[2]
            ])

# ---------- PID & MPC（保留原实现） ----------
class PIDController:
    def __init__(self, kp, ki, kd, max_output=float('inf')):
        self.kp = np.array(kp)
        self.ki = np.array(ki)
        self.kd = np.array(kd)
        self.max_output = max_output
        self.integral = np.zeros(3)
        self.previous_error = np.zeros(3)
        self.previous_time = time.time()
        self.p_term = np.zeros(3)
        self.i_term = np.zeros(3)
        self.d_term = np.zeros(3)
    def compute(self, setpoint, current_value):
        current_time = time.time()
        dt = current_time - self.previous_time
        if dt <= 0:
            return np.zeros(3)
        error = setpoint - current_value
        self.p_term = self.kp * error
        self.integral += error * dt
        max_integral = np.zeros(3)
        for i in range(3):
            if self.ki[i] != 0:
                max_integral[i] = 5.0 / self.ki[i]
            else:
                max_integral[i] = float('inf')
        self.integral = np.clip(self.integral, -max_integral, max_integral)
        self.i_term = self.ki * self.integral
        derivative = (error - self.previous_error) / dt
        self.d_term = self.kd * derivative
        output = self.p_term + self.i_term + self.d_term
        output_norm = np.linalg.norm(output)
        if output_norm > self.max_output:
            output = output / output_norm * self.max_output
        self.previous_error = error
        self.previous_time = current_time
        logging.debug(f"PID - P: {self.p_term}, I: {self.i_term}, D: {self.d_term}, Output: {output}")
        return output

class MPCController:
    def __init__(self, horizon=5, dt=0.1):
        self.horizon = horizon
        self.dt = dt
        self.Q = np.diag([10, 10, 10])
        self.R = np.diag([0.1, 0.1, 0.1])
        self.Qf = np.diag([20, 20, 20])
        self.A = np.eye(3)
        self.B = np.eye(3) * dt
        self.last_control = np.zeros(3)
        self.last_error = np.zeros(3)
    def compute(self, current_state, reference_trajectory):
        current_pos = current_state[:3]
        error = reference_trajectory[0] - current_pos
        Kp = 3.0
        Kd = 2.0
        control = Kp * error + Kd * (self.last_error - error) / self.dt
        control_norm = np.linalg.norm(control)
        if control_norm > 5.0:
            control = control / control_norm * 5.0
        self.last_error = error
        self.last_control = control
        logging.debug(f"MPC - Error: {error}, Control: {control}")
        return control

# ---------- D* Lite Planner（3D 网格的简化实现） ----------
# 说明：这是一个简化但功能齐全的 D*-Lite：在格子上维护 g, rhs, priority queue；
# planner 在 known map 上规划（未知视为可通行），当已知信息变化（新障碍/移除）时调用 update_cell 并增量重规划。
class DStarLitePlanner:
    def __init__(self, grid_size=1.0, bounds=((-50,50),(-50,50),(0,20))):
        self.grid_size = grid_size
        self.bounds = bounds  # 每维 (min,max)
        self.heuristic_cache = {}
        self.rhs = {}  # rhs values
        self.g = {}    # g values
        self.U = []    # heap of (key, node)
        self.entry_finder = {}  # node -> [key, node] for lazy deletion
        self.km = 0.0
        self.start = None  # grid coord (agent)
        self.goal = None   # grid coord (goal)
        self.obstacles = set()  # set of grid coords marked as obstacle (known)
        self.neighbor_offsets = [(dx,dy,dz) for dx in [-1,0,1] for dy in [-1,0,1] for dz in [-1,0,1] if not (dx==0 and dy==0 and dz==0)]
        # cost cache optional
    # utils
    def to_grid(self, pos):
        return (int(round(pos[0] / self.grid_size)), int(round(pos[1] / self.grid_size)), int(round(pos[2] / self.grid_size)))
    def to_world(self, grid):
        return np.array([grid[0]*self.grid_size, grid[1]*self.grid_size, grid[2]*self.grid_size], dtype=float)
    def in_bounds(self, node):
        (minx,maxx),(miny,maxy),(minz,maxz) = self.bounds
        x,y,z = node
        return (minx <= x <= maxx) and (miny <= y <= maxy) and (minz <= z <= maxz)
    def cost(self, a, b):
        # if b is obstacle -> infinite cost
        if b in self.obstacles:
            return float('inf')
        return math.sqrt((a[0]-b[0])**2 + (a[1]-b[1])**2 + (a[2]-b[2])**2)
    def heuristic(self, a, b):
        # Euclidean
        return math.sqrt((a[0]-b[0])**2 + (a[1]-b[1])**2 + (a[2]-b[2])**2)
    def calculate_key(self, node):
        g_val = self.g.get(node, float('inf'))
        rhs_val = self.rhs.get(node, float('inf'))
        val = min(g_val, rhs_val)
        if self.start is None:
            h = 0
        else:
            h = self.heuristic(self.start, node)
        return (val + h + self.km, val)
    def push(self, node):
        key = self.calculate_key(node)
        entry = (key, node)
        heapq.heappush(self.U, entry)
        self.entry_finder[node] = entry
    def remove_entry(self, node):
        # lazy removal: just delete from entry_finder; actual heap entries ignored when popped
        if node in self.entry_finder:
            del self.entry_finder[node]
    def pop(self):
        while self.U:
            key, node = heapq.heappop(self.U)
            # validate
            if node in self.entry_finder and self.entry_finder[node][0] == key:
                del self.entry_finder[node]
                return key, node
            # else stale entry -> skip
        return None, None
    def initialize(self, start_pos, goal_pos):
        # start_pos, goal_pos in world coords
        self.start = self.to_grid(start_pos)
        self.goal = self.to_grid(goal_pos)
        self.rhs.clear()
        self.g.clear()
        self.U.clear()
        self.entry_finder.clear()
        self.km = 0.0
        # initialize
        self.rhs[self.goal] = 0.0
        self.g[self.goal] = float('inf')
        self.push(self.goal)
    def update_vertex(self, u):
        if u != self.goal:
            # rhs(u) = min_{s' in Succ(u)} cost(u,s') + g(s')
            min_rhs = float('inf')
            for dx,dy,dz in self.neighbor_offsets:
                s = (u[0]+dx, u[1]+dy, u[2]+dz)
                if not self.in_bounds(s):
                    continue
                c = self.cost(u, s)
                if c == float('inf'):
                    continue
                g_s = self.g.get(s, float('inf'))
                val = c + g_s
                if val < min_rhs:
                    min_rhs = val
            self.rhs[u] = min_rhs
        # remove old entry and push if inconsistent
        if u in self.entry_finder:
            self.remove_entry(u)
        if self.g.get(u, float('inf')) != self.rhs.get(u, float('inf')):
            self.push(u)
    def compute_shortest_path(self, max_iters=100000):
        iters = 0
        while True:
            iters += 1
            if iters > max_iters:
                logging.warning("D*Lite compute_shortest_path reached max iters")
                break
            top = self.pop()
            if top[0] is None:
                break
            k_old, u = top
            k_new = self.calculate_key(u)
            if k_old < k_new:
                # reinsert with new key
                self.push(u)
            elif self.g.get(u, float('inf')) > self.rhs.get(u, float('inf')):
                self.g[u] = self.rhs[u]
                # update predecessors (neighbors)
                for dx,dy,dz in self.neighbor_offsets:
                    pred = (u[0]+dx, u[1]+dy, u[2]+dz)
                    if not self.in_bounds(pred):
                        continue
                    self.update_vertex(pred)
            else:
                g_old = self.g.get(u, float('inf'))
                self.g[u] = float('inf')
                # update u and predecessors
                self.update_vertex(u)
                for dx,dy,dz in self.neighbor_offsets:
                    pred = (u[0]+dx, u[1]+dy, u[2]+dz)
                    if not self.in_bounds(pred):
                        continue
                    self.update_vertex(pred)
            # termination condition: start is consistent and top key >= key(start)
            start_key = self.calculate_key(self.start)
            if (self.entry_finder.get(self.start) is None and self.g.get(self.start, float('inf')) == self.rhs.get(self.start, float('inf'))):
                break
            # also if top key >= key(start) and start consistent then done (handled by above)
        logging.debug("D*Lite compute finished")
    def update_cell(self, cell, is_obstacle):
        # cell: grid coord; set or clear obstacle and update affected nodes
        old = (cell in self.obstacles)
        if is_obstacle == old:
            return
        if is_obstacle:
            self.obstacles.add(cell)
        else:
            if cell in self.obstacles:
                self.obstacles.remove(cell)
        # when edge cost changes, we increase km (simulates robot movement; simple approach)
        self.km += 0.0  # keep simple; could add distance moved
        # update vertex and neighbors
        self.update_vertex(cell)
        for dx,dy,dz in self.neighbor_offsets:
            n = (cell[0]+dx, cell[1]+dy, cell[2]+dz)
            if self.in_bounds(n):
                self.update_vertex(n)
    def get_shortest_path(self, max_steps=1000):
        # reconstruct path from start following lowest-cost neighbor according to g
        path = []
        current = self.start
        if self.g.get(current, float('inf')) == float('inf'):
            logging.warning("D*Lite: no path from start to goal in known map")
            return []
        path.append(current)
        steps = 0
        while current != self.goal and steps < max_steps:
            steps += 1
            # choose neighbor s that minimizes cost(current,s) + g(s)
            best = None
            best_val = float('inf')
            for dx,dy,dz in self.neighbor_offsets:
                s = (current[0]+dx, current[1]+dy, current[2]+dz)
                if not self.in_bounds(s):
                    continue
                c = self.cost(current, s)
                if c == float('inf'):
                    continue
                val = c + self.g.get(s, float('inf'))
                if val < best_val:
                    best_val = val
                    best = s
            if best is None:
                break
            path.append(best)
            current = best
        # convert to world coords
        world_path = [self.to_world(p) for p in path]
        return world_path

# ---------- Drone（保留，增加传感器参数） ----------
class Drone:
    def __init__(self, x=0, y=0, z=0, mass=1.0):
        self.position = np.array([x, y, z], dtype=float)
        self.velocity = np.array([0, 0, 0], dtype=float)
        self.acceleration = np.array([0, 0, 0], dtype=float)
        self.mass = mass
        self.max_thrust = 30.0
        self.gravity = np.array([0, 0, -9.81])
        self.pid_controller = PIDController(
            kp=np.array([2.0, 2.0, 6.0]),
            ki=np.array([0.1, 0.1, 0.1]),
            kd=np.array([1.0, 1.0, 3.0]),
            max_output=25.0
        )
        self.mpc_controller = MPCController()
        self.controller_type = "PID"
        self.trajectory = []
        self.control_forces = []
        self.step_count = 0
        self.gravity_compensation = np.array([0, 0, 9.81 * self.mass])
        # 传感器参数（感知半径，单位米）
        self.sensor_range = 4.0
    def set_controller(self, controller_type):
        self.controller_type = controller_type
        logging.info(f"切换控制器到: {controller_type}")
    def update(self, target_position, dt=0.1):
        self.step_count += 1
        if self.controller_type == "PID":
            control_force = self.pid_controller.compute(target_position, self.position)
            control_force[2] += self.gravity_compensation[2]
        else:
            ref_trajectory = [target_position] * self.mpc_controller.horizon
            current_state = np.concatenate([self.position, self.velocity])
            control_force = self.mpc_controller.compute(current_state, ref_trajectory)
            control_force[2] += self.gravity_compensation[2]
        self.control_forces.append(control_force.copy())
        total_force = control_force + self.mass * self.gravity
        thrust_norm = np.linalg.norm(total_force)
        if thrust_norm > self.max_thrust:
            total_force = total_force / thrust_norm * self.max_thrust
        self.acceleration = total_force / self.mass
        self.velocity += self.acceleration * dt
        self.position += self.velocity * dt
        if self.position[2] < 0:
            self.position[2] = 0
            self.velocity[2] = 0
        self.trajectory.append(self.position.copy())
        if self.step_count % 10 == 0:
            error = target_position - self.position
            logging.info(
                f"步数 {self.step_count}: 位置={np.round(self.position,2)}, 目标={np.round(target_position,2)}, "
                f"误差={np.round(error,2)}, 控制力={np.round(control_force,2)}"
            )
        distance = np.linalg.norm(target_position - self.position)
        return distance < 0.6  # 松一点容差
    def get_position(self):
        return self.position.copy()
    def get_velocity(self):
        return self.velocity.copy()
    def get_trajectory(self):
        return np.array(self.trajectory)
    def get_control_forces(self):
        return np.array(self.control_forces)

# ---------- NavigationSystem：用 D*Lite 替代 A* 并模拟未知/动态环境 ----------
class NavigationSystem:
    def __init__(self):
        self.log_filename = setup_logging()
        self.data_logger = DataLogger()
        logging.info("初始化导航系统（D*-Lite 版本）...")
        # 真实世界障碍（有些是静态，有些动态）——用于仿真环境
        self.true_obstacles = [
            {'pos': np.array([5, 5, 2]), 'radius': 1.0, 'dynamic': False, 'vel': np.array([0.0,0.0,0.0])},
            {'pos': np.array([8, 3, 3]), 'radius': 1.2, 'dynamic': True,  'vel': np.array([0.05,0.02,0.0])},  # 动态障碍示例
            {'pos': np.array([12, 8, 4]), 'radius': 1.0, 'dynamic': False, 'vel': np.array([0.0,0.0,0.0])},
            {'pos': np.array([15, 12, 5]), 'radius': 1.2, 'dynamic': True,  'vel': np.array([-0.03,0.01,0.0])}, # 动态
            {'pos': np.array([7, 10, 3]), 'radius': 1.0, 'dynamic': False, 'vel': np.array([0.0,0.0,0.0])},
            {'pos': np.array([10, 15, 4]), 'radius': 1.0, 'dynamic': False, 'vel': np.array([0.0,0.0,0.0])},
            {'pos': np.array([13, 5, 3]), 'radius': 1.0, 'dynamic': False, 'vel': np.array([0.0,0.0,0.0])},
            {'pos': np.array([18, 8, 4]), 'radius': 1.0, 'dynamic': False, 'vel': np.array([0.0,0.0,0.0])}
        ]
        # 无人机已知（已观测到）的障碍集合（grid coords）
        self.known_obstacles = set()
        # 起点 终点
        self.start = np.array([0, 0, 1])
        self.goal = np.array([20, 15, 6])
        # planner 使用 D*Lite（在已知地图上规划）
        self.grid_size = 1.0
        self.planner = DStarLitePlanner(grid_size=self.grid_size, bounds=((-50,50),(-50,50),(0,20)))
        # 初始 planner 初始化（start/goal）
        self.planner.initialize(self.start, self.goal)
        # Drone
        self.drone = Drone(x=self.start[0], y=self.start[1], z=self.start[2])
        self.drone.set_controller("PID")
        # 当前路径（世界坐标）和 waypoint index
        self.waypoints = []
        self.current_waypoint_idx = 0
        self.mission_complete = False
        self.start_time = time.time()
        # 初次规划（把已知障碍写入 planner；初始 known_obstacles 为空 -> path 在未知地图上）
        self.replan_from_planner()
        logging.info(f"导航系统初始化完成，起点{self.start} 终点{self.goal}")
    # 将真实世界障碍移动（模拟动态障碍）
    def move_true_obstacles(self, dt=0.1):
        for o in self.true_obstacles:
            if o['dynamic']:
                o['pos'] = o['pos'] + o['vel'] * dt
                # 保持 z 不变且边界内（简单反弹）
                if o['pos'][0] < -40 or o['pos'][0] > 40:
                    o['vel'][0] *= -1
                if o['pos'][1] < -40 or o['pos'][1] > 40:
                    o['vel'][1] *= -1
    # 传感器扫描：基于 drone.sensor_range 去发现真实障碍并更新 known_obstacles -> 更新 planner
    def sensor_scan_and_update(self):
        sensor_range = self.drone.sensor_range
        drone_pos = self.drone.get_position()
        newly_detected = []
        for o in self.true_obstacles:
            # if center within sensor_range + radius -> detect
            dist = np.linalg.norm(o['pos'] - drone_pos)
            if dist <= sensor_range + o['radius']:
                # 将障碍体积映射到多个网格并加入 known_obstacles
                r_cells = int(math.ceil(o['radius'] / self.grid_size))
                center_grid = self.planner.to_grid(o['pos'])
                for dx in range(-r_cells, r_cells+1):
                    for dy in range(-r_cells, r_cells+1):
                        for dz in range(-r_cells, r_cells+1):
                            cell = (center_grid[0]+dx, center_grid[1]+dy, center_grid[2]+dz)
                            if not self.planner.in_bounds(cell):
                                continue
                            # optional: 仅加入球内格子
                            if dx*dx + dy*dy + dz*dz <= (r_cells+0.5)**2:
                                if cell not in self.known_obstacles:
                                    self.known_obstacles.add(cell)
                                    # update planner incrementally
                                    self.planner.update_cell(cell, is_obstacle=True)
                                    newly_detected.append(cell)
        if newly_detected:
            logging.info(f"传感器发现 {len(newly_detected)} 个新障碍（格子）。触发重规划。")
            self.replan_from_planner()
    # 当已知障碍增加或变化时从 planner 获取新路径
    def replan_from_planner(self):
        # compute shortest path on known map
        self.planner.start = self.planner.to_grid(self.drone.get_position())
        # ensure start is in planner structures
        if self.planner.start not in self.planner.g:
            self.planner.g[self.planner.start] = float('inf')
            self.planner.rhs[self.planner.start] = float('inf')
        self.planner.compute_shortest_path()
        raw_path = self.planner.get_shortest_path()
        if not raw_path:
            logging.warning("重规划未找到路径（已知地图），尝试返回直线路径")
            self.waypoints = [self.goal]
            self.current_waypoint_idx = 0
        else:
            # path smoothing（简单地去除连续共线点）
            smoothed = self.simple_smooth(raw_path)
            self.waypoints = smoothed
            self.current_waypoint_idx = 0
            logging.info(f"重规划完成，得到 {len(self.waypoints)} 个航点。")
    def simple_smooth(self, path):
        # 去除冗余共线点（保留端点）
        if len(path) <= 2:
            return path
        res = [path[0]]
        for i in range(1, len(path)-1):
            a = np.array(res[-1])
            b = np.array(path[i])
            c = np.array(path[i+1])
            # 如果 b 在 a->c 的线性近似上（共线）则跳过
            ab = b - a
            ac = c - a
            if np.linalg.norm(np.cross(ab, ac)) < 1e-6:
                continue
            res.append(path[i])
        res.append(path[-1])
        return res
    def plan_path_initial(self):
        # 适配接口：初始从 planner 获取 path
        self.replan_from_planner()
    # 更新流程：移动真障碍 -> 传感器扫描 -> 驱动无人机跟踪当前 waypoint（若障碍出现则 replan）
    def update(self, dt=0.1):
        if self.mission_complete:
            return True
        # move true dynamic obstacles
        self.move_true_obstacles(dt)
        # sensor scan (drone perceives some obstacles)
        self.sensor_scan_and_update()
        # get current target waypoint
        if self.current_waypoint_idx >= len(self.waypoints):
            target = self.goal.copy()
        else:
            target = self.waypoints[self.current_waypoint_idx]
        # update drone
        reached = self.drone.update(target, dt)
        # log
        current_time = time.time() - self.start_time
        self.data_logger.log_data(
            current_time,
            self.drone.get_position(),
            self.drone.get_velocity(),
            np.array(target),
            self.drone.control_forces[-1] if self.drone.control_forces else np.zeros(3)
        )
        if reached:
            logging.info(f"到达航点 {self.current_waypoint_idx}: {target}")
            self.current_waypoint_idx += 1
            if self.current_waypoint_idx >= len(self.waypoints):
                logging.info("所有航点完成或到达最终目标！")
                self.mission_complete = True
                return True
        # Additional check: if next waypoint lies inside known obstacle, replan
        if self.current_waypoint_idx < len(self.waypoints):
            next_wp = self.waypoints[self.current_waypoint_idx]
            grid_wp = self.planner.to_grid(next_wp)
            if grid_wp in self.known_obstacles:
                logging.info("下一个航点发生障碍阻塞，触发重规划")
                self.replan_from_planner()
        return False
    # helpers for viz
    def get_drone_position(self):
        return self.drone.get_position()
    def get_waypoints(self):
        return self.waypoints
    def get_trajectory(self):
        return self.drone.get_trajectory()
    def get_control_forces(self):
        return self.drone.get_control_forces()
    def get_true_obstacles(self):
        return self.true_obstacles
    def get_known_obstacle_grids(self):
        return self.known_obstacles
    def set_controller(self, controller_type):
        self.drone.set_controller(controller_type)
    def print_summary(self):
        trajectory = self.get_trajectory()
        total_distance = 0
        for i in range(1, len(trajectory)):
            total_distance += np.linalg.norm(trajectory[i] - trajectory[i - 1])
        logging.info(f"任务摘要: 总飞行距离={total_distance:.2f}m, 航点数={len(self.waypoints)}")
        logging.info(f"日志文件: {self.log_filename}")

# ---------- Simulation & Visualization（3D） ----------
class Simulation:
    def __init__(self, controller_type="PID"):
        self.nav_system = NavigationSystem()
        self.nav_system.set_controller(controller_type)
        self.fig = plt.figure(figsize=(15,10))
        self.ax = self.fig.add_subplot(111, projection='3d')
        self.ax.set_title(f'无人机导航模拟 - D*-Lite（未知/动态障碍） - {controller_type}')
        self.ax.set_xlabel('X (m)')
        self.ax.set_ylabel('Y (m)')
        self.ax.set_zlabel('Z (m)')
        self.ax.set_xlim(-2,25)
        self.ax.set_ylim(-2,20)
        self.ax.set_zlim(0,10)
        start = self.nav_system.start
        goal = self.nav_system.goal
        self.ax.scatter(start[0], start[1], start[2], c='green', marker='o', s=100, label='起点')
        self.ax.scatter(goal[0], goal[1], goal[2], c='red', marker='o', s=100, label='终点')
        # 绘制已知/未知障碍分别用不同透明度
        self.true_obs_artists = []
        self.known_obs_artists = []
        self.waypoint_artist = None
        # 初始绘制真实障碍（灰色，透明）
        for o in self.nav_system.get_true_obstacles():
            self.draw_cube(o['pos'], o['radius']*2.0, color='gray', alpha=0.25)
        # 绘制已知障碍（更明显）
        for cell in self.nav_system.get_known_obstacle_grids():
            world = np.array(self.nav_system.planner.to_world(cell))
            self.draw_cube(world, 1.0, color='black', alpha=0.6)
        # path (规划路径)
        self.path_line, = self.ax.plot([],[],[], 'r--', alpha=0.6, linewidth=2, label='规划路径（已知地图）')
        # drone marker & trajectory
        self.drone_marker, = self.ax.plot([],[],[], 'bo', markersize=8, label='无人机')
        self.trajectory_line, = self.ax.plot([],[],[], 'b-', alpha=0.7, linewidth=2, label='实际轨迹')
        self.target_marker, = self.ax.plot([],[],[], 'rx', markersize=10, label='当前目标')
        self.ax.legend()
        self.ax.view_init(elev=20, azim=45)
        self.info_text = self.ax.text2D(0.02, 0.95, "", transform=self.ax.transAxes)
        self.animation = None
        self.frame_count = 0
    def draw_cube(self, center, size, color='gray', alpha=0.25):
        r = size / 2.0
        vertices = [
            [center[0] - r, center[1] - r, center[2] - r],
            [center[0] + r, center[1] - r, center[2] - r],
            [center[0] + r, center[1] + r, center[2] - r],
            [center[0] - r, center[1] + r, center[2] - r],
            [center[0] - r, center[1] - r, center[2] + r],
            [center[0] + r, center[1] - r, center[2] + r],
            [center[0] + r, center[1] + r, center[2] + r],
            [center[0] - r, center[1] + r, center[2] + r]
        ]
        faces = [
            [0,1,2,3],[4,5,6,7],[0,1,5,4],[2,3,7,6],[0,3,7,4],[1,2,6,5]
        ]
        poly_verts = []
        for face in faces:
            poly_verts.append([vertices[i] for i in face])
        collection = Poly3DCollection(poly_verts, facecolors=color, linewidths=0.5, edgecolors='k', alpha=alpha)
        self.ax.add_collection3d(collection)
        return collection
    def update_plot(self, frame):
        self.frame_count += 1
        complete = self.nav_system.update(dt=0.1)
        drone_pos = self.nav_system.get_drone_position()
        # update true obstacles: redraw by clearing & redrawing is simpler - here we just update by clearing axes children and replot everything each frame (成本可接受)
        self.ax.collections.clear()
        # redraw static markers
        start = self.nav_system.start
        goal = self.nav_system.goal
        self.ax.scatter(start[0], start[1], start[2], c='green', marker='o', s=100)
        self.ax.scatter(goal[0], goal[1], goal[2], c='red', marker='o', s=100)
        # draw true obstacles (light)
        for o in self.nav_system.get_true_obstacles():
            self.draw_cube(o['pos'], o['radius']*2.0, color='gray', alpha=0.25)
        # draw known obstacles (solid)
        for cell in self.nav_system.get_known_obstacle_grids():
            world = np.array(self.nav_system.planner.to_world(cell))
            self.draw_cube(world, 1.0, color='black', alpha=0.6)
        # draw planned path
        waypoints = self.nav_system.get_waypoints()
        if waypoints:
            wp_x = [wp[0] for wp in waypoints]
            wp_y = [wp[1] for wp in waypoints]
            wp_z = [wp[2] for wp in waypoints]
            self.path_line.set_data(wp_x, wp_y)
            self.path_line.set_3d_properties(wp_z)
        else:
            self.path_line.set_data([],[])
            self.path_line.set_3d_properties([])
        # update drone marker & trajectory
        self.drone_marker.set_data([drone_pos[0]], [drone_pos[1]])
        self.drone_marker.set_3d_properties([drone_pos[2]])
        trajectory = self.nav_system.get_trajectory()
        if len(trajectory) > 1:
            self.trajectory_line.set_data(trajectory[:,0], trajectory[:,1])
            self.trajectory_line.set_3d_properties(trajectory[:,2])
        # current target
        if self.nav_system.current_waypoint_idx < len(self.nav_system.waypoints):
            target_pos = self.nav_system.waypoints[self.nav_system.current_waypoint_idx]
        else:
            target_pos = self.nav_system.goal
        self.target_marker.set_data([target_pos[0]], [target_pos[1]])
        self.target_marker.set_3d_properties([target_pos[2]])
        info = f"帧: {self.frame_count}\n位置: {np.round(drone_pos,2)}\n目标: {np.round(target_pos,2)}\n已知障碍格数: {len(self.nav_system.get_known_obstacle_grids())}"
        if complete:
            info += "\n状态: 任务完成!"
        self.info_text.set_text(info)
        if complete:
            self.animation.event_source.stop()
            self.nav_system.print_summary()
            logging.info("模拟完成")
        return self.drone_marker, self.trajectory_line, self.path_line, self.target_marker, self.info_text
    def run(self):
        logging.info("开始 D*-Lite 无人机导航模拟...")
        self.animation = FuncAnimation(self.fig, self.update_plot, frames=2000, interval=80, blit=False, repeat=False)
        plt.show()


# ---------- main ----------
if __name__ == "__main__":
    sim = Simulation(controller_type="PID")
    sim.run()
